import os
import csv

from data.deal_filename import remove_spaces


def list_files(folder_path):
    # 获取文件夹及其子目录下的所有文件
    files = []
    sub_string = '.DS_Store'  # 剔除苹果的系统文件
    for root, dirs, filenames in os.walk(folder_path):
        for filename in filenames:
            if sub_string not in filename:
                files.append(os.path.join(root, filename))
    return files


def write_to_csv(folder_path, csv_file, start_num):
    # 获取文件夹及其子目录下的所有文件
    files = list_files(folder_path)

    # 打开 CSV 文件进行写入
    with open(csv_file, 'w', newline='') as csvfile:
        # 定义 CSV 文件的表头
        fieldnames = ['id', 'path']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

        # 写入表头
        writer.writeheader()

        # 遍历文件列表，写入每一行数据
        for idx, file_path in enumerate(files, start=start_num):
            writer.writerow({'id': idx, 'path': "./" + file_path})


if __name__ == "__main__":
    folder_path = 'test_AI测试图片/1221-1'  # 替换为你的文件夹路径
    csv_file = f'csv/test_AI测试图片_1221-1.csv'  # 输出的 CSV 文件名
    start_num = 9690

    # 先做文件名的修正
    remove_spaces(folder_path)
    # 调用函数写入 CSV 文件
    write_to_csv(folder_path, csv_file, start_num)
